import numpy as np
import numba as nb
from numba import cuda
from math import exp, ceil
from ..utils import get_grid_1d, get_grid_2d, assume_ndarray, flatten


@cuda.jit()
def _gpu_multiply_elements(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    """对应位置处元素相乘"""
    x = cuda.grid(1)
    if x < a.shape[0]:
        c[x] = a[x] * b[x]


@cuda.reduce
def _gpu_reduce_sum(a, b):
    return a + b


@cuda.jit()
def _gpu_matmul_2d(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    """二维矩阵乘法"""
    x, y = cuda.grid(2)
    if x < c.shape[0] and y < c.shape[1]:
        sum = 0
        for i in range(a.shape[1]):
            sum += a[x, i] * b[i, y]
        c[x, y] = sum


@cuda.jit()
def _gpu_matmul_0d(a: np.ndarray, b, c: np.ndarray):
    """矩阵乘以标量"""
    x = cuda.grid(1)
    if x < a.shape[0]:
        c[x] = a[x] * b


@cuda.jit()
def _gpu_matmul_1d(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    """二维矩阵与一维向量乘法"""
    x = cuda.grid(1)
    if x < c.shape[0]:
        sum = 0
        for i in range(a.shape[1]):
            sum += a[x, i] * b[i]
        c[x] = sum


def dot(a: np.ndarray, b: np.ndarray):
    a, b = assume_ndarray(a, b)

    if a.ndim == 1 and b.ndim == 1:
        assert a.shape[0] == b.shape[0]
        c = np.empty_like(a)

        grid, block = get_grid_1d(a.shape[0])
        _gpu_multiply_elements[grid, block](a, b, c)

        return _gpu_reduce_sum(c)
    elif a.ndim == 2 and b.ndim == 2:
        assert a.shape[1] == b.shape[0]
        c = np.empty((a.shape[0], b.shape[1]))

        grid, block = get_grid_2d(c.shape[0], c.shape[1])
        _gpu_matmul_2d[grid, block](a, b, c)

        return c
    elif a.ndim == 0 or b.ndim == 0:
        if a.ndim == 0 and b.ndim == 0:
            return a * b

        if a.ndim == 0:
            a, b = b, a

        b = b.item()
        c = np.empty_like(a)

        grid, block = get_grid_1d(c.size)
        _gpu_matmul_0d[grid, block](flatten(a), b, flatten(c))

        return c
    elif b.ndim == 1:
        assert a.shape[-1] == b.shape[0]
        shape = a.shape[:-1]

        a = a.reshape(-1, a.shape[-1])
        c = np.empty((a.shape[0],))

        grid, block = get_grid_1d(a.shape[0])
        _gpu_matmul_1d[grid, block](a, b, c)

        return c.reshape(shape)
    elif b.ndim >= 2:
        assert a.shape[-1] == b.shape[0]
        shape = a.shape[:-1] + b.shape[1:]

        a = a.reshape((-1, a.shape[-1]))
        b = b.reshape((b.shape[0], -1))
        c = np.empty((a.shape[0], b.shape[1]))

        grid, block = get_grid_2d(c.shape[0], c.shape[1])
        _gpu_matmul_2d[grid, block](a, b, c)

        return c.reshape(shape)
    else:
        raise Exception("参数错误")
